import os
import torch
import model.data_loader as data_loader
import model.net as nn
import model.net_2 as nn2
import utils
from torch.autograd import Variable
import torch.nn.functional as F
import  math
import numpy as np
import tool as tl

def get_dis(path):
    path = path + '/result/'
    files = os.listdir(path)
    dis_list = []
    for file in files:
        path1 = path + file
        dis_list.append(tl.to_float(tl.read_csv(path1))[0][0])
    return np.mean(dis_list)


def get_ie(model, checkpoint, params, val_dataloader, device):
    utils.load_checkpoint(checkpoint, model)
    model.eval()
    # summary for current eval loop
    summ = 0
    for data_batch, labels_batch in val_dataloader:
        data_batch, labels_batch = Variable(data_batch), Variable(labels_batch)
        # compute model output
        with torch.no_grad():
            output_teacher_batch = model(data_batch)
        out = F.softmax(output_teacher_batch)
        for idx, label in enumerate(labels_batch):
            # print(out[idx])
            tmp = out[idx]
            for p in out[idx]:
                #if p>0.999999 or p<0.000001:
                   # continue
                #else:
                summ += p*math.log(p)
    return summ

if __name__=='__main__':
    torch.cuda.set_device(1)
    device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")

    json_path = 'experiments/base_cnn/params.json'
    params = params = utils.Params(json_path)
    params.cuda = torch.cuda.is_available()
    print('the gpu is {}'.format(params.cuda))
    print('the current gpu is {}'.format(torch.cuda.current_device()))
    model = nn2.Net(params)

    dev_dl = data_loader.fetch_dataloader('dev', params)

    root = 'epoch_params/60epoch/'
    distance_root = 'weights_params/IE demo val loss 60 epoch/'
    baseline_checkpoint = 'epoch_params/160epoch/baseline 2layer cnn/2layercnn159/best.pth.tar'
    baseline_ie = get_ie(model, baseline_checkpoint, params, dev_dl, device)
    baseline_checkpoint = 'epoch_params/160epoch/baseline 2layer cnn/2layercnn-1/best.pth.tar'
    baseline_ie_init = get_ie(model, baseline_checkpoint, params, dev_dl, device)
    baseline_ie = baseline_ie_init - baseline_ie
    baseline_dis = get_dis(distance_root + 'baseline 2layer cnn/')

    IE_res = []
    IE_class = []

    class_name = os.listdir(root)
    for name in class_name:
        if 'layer_' in name:
            print(name)
            IE_class.append(name)
            ave_paths = root + name + '/'
            turns = os.listdir(ave_paths)
            KD_ie = []
            KD_ie_init = []
            for turn in turns:
                KD_checkpoint = ave_paths + turn + '/159/best.pth.tar'
                KD_ie.append(get_ie(model, KD_checkpoint, params, dev_dl, device))
                KD_checkpoint = ave_paths + turn + '/-1/best.pth.tar'
                KD_ie_init.append(get_ie(model, KD_checkpoint, params, dev_dl, device))
            KD_ie_ave = np.mean(np.array(KD_ie))
            KD_ie_ave_init = np.mean(np.array(KD_ie_init))
            KD_ie_ave = KD_ie_ave_init - KD_ie_ave
            KD_dis_ave = get_dis(distance_root + name + '/')
            print((KD_ie_ave * math.log(KD_dis_ave))/ (baseline_ie * math.log(baseline_dis)))
            IE_res.append((KD_ie_ave * math.log(KD_dis_ave))/ (baseline_ie * math.log(baseline_dis)))
            # print((KD_ie_ave) / (baseline_ie))
            # IE_res.append((KD_ie_ave) / (baseline_ie))
    for i in range(len(IE_res)):
        print('{}: {}'.format(IE_class[i], IE_res[i]))

